import tensorflow as tf
import collections
import matplotlib.pyplot as plt
import sys
def read(filename):
    data = collections.defaultdict(dict)
    for event in tf.compat.v1.train.summary_iterator(filename):
        for value in event.summary.value:
            data[value.tag][event.step] = tf.make_ndarray(value.tensor)
    return data
def plot(ds, tags, directory):
    fig, ax = plt.subplots()
    ax.set_title(" and ".join(tags))
    ax.set_xlabel('steps')
    for k,d in ds.items():
        for tag in tags:
            data = d[tag]
            x = sorted(list(data.keys()))
            y = [data[key] for key in x]
            ax.plot(x,y,marker=".",markersize=1,label=k + " " + tag)
    ax.legend()
    filename = "_and_".join(tags).replace("/",".")
    fig.savefig(F'{directory}/{filename}.png')
    return
    
if __name__=="__main__":
    print(sys.argv)
    if len(sys.argv) == 6:
        filename1 = sys.argv[1]
        desc1 = sys.argv[2]
        filename2 = sys.argv[3]
        desc2 = sys.argv[4]
        ds = dict()
        ds[desc1] = read(filename1)
        ds[desc2] = read(filename2)
        plot(ds, ["eval/episode_reward"], sys.argv[5])
    elif len(sys.argv) == 3:
        filename = sys.argv[1]
        ds = { " ": read(filename)}
        plot(ds, ["eval/episode_reward"], sys.argv[2])
